package chan.project.manager;

import com.baidubce.qianfan.Qianfan;
import com.baidubce.qianfan.core.StreamIterator;
import com.baidubce.qianfan.model.chat.ChatRequest;
import com.baidubce.qianfan.model.chat.ChatResponse;
import com.baidubce.qianfan.model.chat.Message;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;

@Component
public class QianFanManager {

    @Resource
    private Qianfan qianfan;

    /**
     * 简化普通调用(只调用一次)
     * @param message
     * @return
     */
    public String Chat(String message){
        Message sendMessage = new Message();
        sendMessage.setContent(message);
        sendMessage.setRole("user");
        List<Message> messages = new ArrayList<>();
        messages.add(sendMessage);
        return Chat(messages);
    }

    /**
     * 普通调用（多次调用）
     * @param messageList
     * @return
     */
    public String Chat(String[] messageList){
        List<Message> messages = new ArrayList<>();
        for (String message : messageList) {
            Message sendMessage = new Message();
            sendMessage.setContent(message);
            sendMessage.setRole("user");
            messages.add(sendMessage);
        }
        return Chat(messages);
    }

    /**
     * 原始普通调用
     * @param messages
     * @return
     */
    public String Chat(List<Message> messages){
//        messages中最后一个是当前请求信息，之前的代表历史请求信息
//        奇数位的role要是user，偶数位的为assistant,并且messages的长度必定要为奇数，否则无法执行
        ChatResponse response = qianfan.chatCompletion()
                .model("ERNIE-Speed-8K") // 使用model指定预置模型
                .messages(messages) // 添加用户消息 (此方法可以调用多次，以实现多轮对话的消息传递)
                .temperature(0.7) // 文本随机数
                .execute(); // 发起请求
        return response.getResult();
    }


    /**
     * 简化流式调用(只调用一次)
     * @param message
     * @return
     */
    public StreamIterator<ChatResponse> ChatByStream(String message){
        Message sendMessage = new Message();
        sendMessage.setContent(message);
        sendMessage.setRole("user");
        List<Message> messages = new ArrayList<>();
        messages.add(sendMessage);
        return ChatByStream(messages);
    }

    /**
     * 流式调用（多次调用）
     * @param messageList
     * @return
     */
    public StreamIterator<ChatResponse> ChatByStream(String[] messageList){
        List<Message> messages = new ArrayList<>();
        for (String message : messageList) {
            Message sendMessage = new Message();
            sendMessage.setContent(message);
            sendMessage.setRole("user");
            messages.add(sendMessage);
        }
        return ChatByStream(messages);
    }

    /**
     * 流式接口的调用(返回调用对象)
     * @param messages
     * @return
     */
    public StreamIterator<ChatResponse> ChatByStream(List<Message> messages){
        ChatRequest chatRequest = new ChatRequest();
        chatRequest.setMessages(messages);
        chatRequest.setStream(true);
        chatRequest.setModel("ERNIE-3.5-8K");
        chatRequest.setTemperature(0.7);
        StreamIterator<ChatResponse> chatResponseStreamIterator = qianfan.chatCompletionStream(chatRequest);
        return chatResponseStreamIterator;
    }
}
